from pathlib import Path
import random
import cv2
import os

def reorganize(cls,input_dir,output_dir):
    all_files = []
    for file in input_dir.glob("*.jpg"):
        all_files.append(file)
    for file in input_dir.glob("*.png"):
        all_files.append(file)
    num_imgs = len(all_files)
    val_files,train_files=[],[]
    for file in random.sample(all_files,int(0.1*num_imgs)):
        val_files.append(str(file))
    for file in all_files:
        if file not in val_files:
            train_files.append(file)

    output_val_dir = output_dir.format("val",cls)
    output_train_dir = output_dir.format("train",cls)
    for dir in [output_val_dir,output_train_dir]:
        if not os.path.exists(dir):
            os.makedirs(dir)

    for file in train_files:
         img_data= cv2.imread(str(file))
         # print(output_dir+Path(file).stem+".jpg")
         cv2.imwrite(output_train_dir+str(Path(file).stem)+".jpg",img_data)

    for file in val_files:
        img_data = cv2.imread(str(file))
        # print(output_dir+Path(file).stem+".jpg")
        cv2.imwrite(output_val_dir + str(Path(file).stem) + ".jpg", img_data)

if __name__=="__main__":
    output_dir="/mnt/c/JXLZ/医疗发票/classifier/trainning_data/{}/{}/"
    bejing_img_dir = Path("/mnt/c/JXLZ/医疗发票/classifier/北京票据/default")
    shanghai_img_dir = Path("/mnt/c/JXLZ/医疗发票/classifier/上海票据/default")
    reorganize("shanghai",shanghai_img_dir,output_dir)

    # for file in bejing_train_files:
    #      img_data= cv2.imread(file)
    #      # print(output_dir+Path(file).stem+".jpg")
    #      cv2.imwrite(output_dir+str(Path(file).stem)+".jpg",img_data)
    # # pth = "/mnt/c/JXLZ/医疗发票/classifier/上海票据/default/med_inv_23.jpg"
    # # print(cv2.imread(pth))






